# -*- coding: utf-8 -*-
"""
13日均线在21日均线之上筛选器
筛选规则：
1. 最近连续20个交易日，13日均线都在21日均线之上
2. 显示最新的均线数据和价格信息

这种均线排列的情况可能预示着：
1. 股票处于中期上升趋势
2. 短期均线强于中期均线
3. 趋势保持稳定且持续时间长
4. 潜在的趋势延续机会

使用方法：
    python ma10_above_ma20_screener.py              # 筛选所有股票
    python ma10_above_ma20_screener.py --limit 100  # 限制筛选数量（测试用）
    python ma10_above_ma20_screener.py --stock 000001  # 筛选指定股票
"""
import logging
import pandas as pd
import numpy as np
from datetime import datetime
import sys
import os
import argparse

# 添加父目录到路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from database import DatabaseManager, StockInfoDAO
from config import DATABASE_CONFIG
import warnings

# 忽略pandas警告
warnings.filterwarnings('ignore', category=pd.errors.SettingWithCopyWarning)

# 配置日志
# 获取项目根目录
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
LOG_DIR = os.path.join(PROJECT_ROOT, 'logs')
os.makedirs(LOG_DIR, exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler(os.path.join(LOG_DIR, 'ma13_above_ma21_screener.log'), encoding='utf-8')
    ]
)
logger = logging.getLogger(__name__)


class MA10AboveMA20Screener:
    """13日均线在21日均线之上筛选器"""
    
    def __init__(self):
        self.db = DatabaseManager(DATABASE_CONFIG)
        self.stock_dao = StockInfoDAO(self.db)
    
    def get_stock_codes(self, limit: int = None) -> list:
        """获取所有活跃股票代码"""
        try:
            sql = "SELECT stock_code, stock_name FROM new_stock_info WHERE status = 'active' ORDER BY stock_code"
            if limit:
                sql += f" LIMIT {limit}"
            
            results = self.db.execute_query(sql)
            return [(row['stock_code'], row['stock_name']) for row in results]
        except Exception as e:
            logger.error(f"获取股票列表失败: {e}")
            return []
    
    def get_stock_daily_data(self, stock_code: str, days: int = 60) -> pd.DataFrame:
        """
        获取股票最近N天的日线数据
        
        Args:
            stock_code: 股票代码
            days: 获取天数（默认60天，确保有足够的数据计算21日均线+20天检查）
            
        Returns:
            包含交易数据的DataFrame
        """
        sql = """
        SELECT trade_date, close_price, change_pct, volume, amount
        FROM new_stock_daily 
        WHERE stock_code = %s 
        ORDER BY trade_date DESC 
        LIMIT %s
        """
        try:
            results = self.db.execute_query(sql, (stock_code, days))
            if not results:
                return pd.DataFrame()
            
            df = pd.DataFrame(results)
            df['trade_date'] = pd.to_datetime(df['trade_date'])
            
            # 转换数值列
            df['close_price'] = pd.to_numeric(df['close_price'], errors='coerce')
            df['change_pct'] = pd.to_numeric(df['change_pct'], errors='coerce')
            df['volume'] = pd.to_numeric(df['volume'], errors='coerce')
            df['amount'] = pd.to_numeric(df['amount'], errors='coerce')
            
            # 按日期升序排列
            df = df.sort_values('trade_date').reset_index(drop=True)
            return df
            
        except Exception as e:
            logger.error(f"获取股票 {stock_code} 数据失败: {e}")
            return pd.DataFrame()
    
    def calculate_moving_averages(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        计算移动平均线
        
        Args:
            df: 股票日线数据DataFrame
            
        Returns:
            添加了均线列的DataFrame
        """
        if df.empty:
            return df
        
        # 计算各周期均线
        df['ma5'] = df['close_price'].rolling(window=5).mean()
        df['ma13'] = df['close_price'].rolling(window=13).mean()
        df['ma21'] = df['close_price'].rolling(window=21).mean()
        df['ma30'] = df['close_price'].rolling(window=30).mean()
        
        return df
    
    def check_ma10_above_ma20(self, df: pd.DataFrame, check_days: int = 20) -> dict:
        """
        检查是否满足13日均线在21日均线之上的条件
        
        条件：
        1. 最近连续N个交易日（默认20天），13日均线都在21日均线之上
        
        Args:
            df: 股票日线数据DataFrame（已计算均线）
            check_days: 检查的交易日数量（默认20天）
            
        Returns:
            符合条件返回详细信息字典，否则返回None
        """
        if df.empty or len(df) < 21:  # 至少需要21天数据来计算21日均线
            return None
        
        # 获取最近N个交易日的数据
        recent_days = df.tail(check_days).copy()
        
        if len(recent_days) < check_days:
            return None
        
        # 检查最近N天的数据是否有空值（只检查我们关心的这几天）
        if recent_days[['ma13', 'ma21']].isna().any().any():
            return None
        
        # 条件1: 最近连续N个交易日，13日均线都在21日均线之上
        ma10_above_ma20 = (recent_days['ma13'] > recent_days['ma21']).all()
        
        if not ma10_above_ma20:
            return None
        
        # 获取最后一个交易日的数据
        latest_row = df.iloc[-1]
        
        # 计算额外的统计信息
        # 计算均线之间的距离（作为趋势强度指标）
        ma10_ma20_gap = ((latest_row['ma13'] - latest_row['ma21']) / latest_row['ma21'] * 100)
        
        # 计算最近N天的平均间距
        recent_days['ma10_ma20_gap'] = ((recent_days['ma13'] - recent_days['ma21']) / recent_days['ma21'] * 100)
        avg_gap = recent_days['ma10_ma20_gap'].mean()
        min_gap = recent_days['ma10_ma20_gap'].min()
        max_gap = recent_days['ma10_ma20_gap'].max()
        
        # 计算最近N天的价格变化
        first_close = recent_days.iloc[0]['close_price']
        last_close = recent_days.iloc[-1]['close_price']
        total_change = ((last_close - first_close) / first_close * 100)
        
        # 计算最近N天的平均涨幅
        avg_change = recent_days['change_pct'].mean()
        max_gain = recent_days['change_pct'].max()
        max_loss = recent_days['change_pct'].min()
        
        # 计算最近N天的平均成交量和成交额
        avg_volume = recent_days['volume'].mean()
        avg_amount = recent_days['amount'].mean()
        
        # 计算均线斜率（趋势方向），使用最近5天的数据
        if len(df) >= 5:
            ma10_slope = ((df['ma13'].iloc[-1] - df['ma13'].iloc[-5]) / df['ma13'].iloc[-5] * 100)
            ma20_slope = ((df['ma21'].iloc[-1] - df['ma21'].iloc[-5]) / df['ma21'].iloc[-5] * 100)
        else:
            ma10_slope = 0
            ma20_slope = 0
        
        return {
            'trade_date': latest_row['trade_date'],
            'close_price': latest_row['close_price'],
            'ma5': round(latest_row['ma5'], 2) if pd.notna(latest_row['ma5']) else None,
            'ma13': round(latest_row['ma13'], 2),
            'ma21': round(latest_row['ma21'], 2),
            'ma30': round(latest_row['ma30'], 2) if pd.notna(latest_row['ma30']) else None,
            'ma13_ma21_gap': round(ma10_ma20_gap, 2),
            'avg_gap': round(avg_gap, 2),
            'min_gap': round(min_gap, 2),
            'max_gap': round(max_gap, 2),
            'ma13_slope': round(ma10_slope, 2),
            'ma21_slope': round(ma20_slope, 2),
            'total_change': round(total_change, 2),
            'avg_change': round(avg_change, 2),
            'max_gain': round(max_gain, 2),
            'max_loss': round(max_loss, 2),
            'latest_change_pct': round(latest_row['change_pct'], 2) if pd.notna(latest_row['change_pct']) else 0,
            'avg_volume': round(avg_volume, 2),
            'avg_amount': round(avg_amount, 2),
            'check_days': check_days
        }
    
    def screen_stocks(self, stock_codes: list = None, limit: int = None, check_days: int = 20) -> pd.DataFrame:
        """
        筛选符合条件的股票
        
        Args:
            stock_codes: 指定要筛选的股票列表，格式为[(code, name), ...]
            limit: 限制筛选的股票数量
            check_days: 检查的交易日数量（默认20天）
            
        Returns:
            包含筛选结果的DataFrame
        """
        if stock_codes is None:
            stock_codes = self.get_stock_codes(limit=limit)
        
        if not stock_codes:
            logger.warning("没有找到需要筛选的股票")
            return pd.DataFrame()
        
        logger.info(f"开始筛选 {len(stock_codes)} 只股票（检查最近{check_days}天）...")
        
        results = []
        total = len(stock_codes)
        
        for idx, (stock_code, stock_name) in enumerate(stock_codes, 1):
            try:
                # 每10只股票输出一次进度
                if idx % 10 == 0 or idx == total:
                    logger.info(f"进度: {idx}/{total} ({idx/total*100:.1f}%)")
                
                # 获取日线数据（需要更多数据来计算均线）
                df = self.get_stock_daily_data(stock_code, days=60)
                
                if df.empty or len(df) < 21:
                    continue
                
                # 计算均线
                df = self.calculate_moving_averages(df)
                
                # 检查是否满足条件
                result = self.check_ma10_above_ma20(df, check_days=check_days)
                
                if result:
                    result['stock_code'] = stock_code
                    result['stock_name'] = stock_name
                    results.append(result)
                    
                    logger.info(
                        f"✓ {stock_code} {stock_name} - "
                        f"价格: {result['close_price']:.2f}, "
                        f"MA13: {result['ma13']:.2f}, MA21: {result['ma21']:.2f}, "
                        f"间距: {result['ma13_ma21_gap']:.2f}%, "
                        f"{check_days}日涨幅: {result['total_change']:.2f}%"
                    )
            
            except Exception as e:
                logger.error(f"处理股票 {stock_code} {stock_name} 时出错: {e}")
                continue
        
        if not results:
            logger.info("未找到符合条件的股票")
            return pd.DataFrame()
        
        # 转换为DataFrame并排序
        df_results = pd.DataFrame(results)
        
        # 按MA13-MA21间距降序排列（间距越大说明短期趋势越强）
        df_results = df_results.sort_values('ma13_ma21_gap', ascending=False)
        
        # 重新排列列的顺序
        columns_order = [
            'stock_code', 'stock_name', 'trade_date', 'close_price',
            'ma5', 'ma13', 'ma21', 'ma30',
            'ma13_ma21_gap', 'avg_gap', 'min_gap', 'max_gap',
            'ma13_slope', 'ma21_slope',
            'latest_change_pct', 'total_change', 'avg_change',
            'max_gain', 'max_loss',
            'avg_volume', 'avg_amount', 'check_days'
        ]
        df_results = df_results[columns_order]
        
        logger.info(f"\n筛选完成！共找到 {len(df_results)} 只符合条件的股票")
        
        return df_results
    
    def save_results(self, df: pd.DataFrame, output_dir: str = '.') -> tuple:
        """
        保存筛选结果到CSV和TXT文件
        
        Args:
            df: 筛选结果DataFrame
            output_dir: 输出目录
            
        Returns:
            (csv_file, txt_file) 文件路径元组
        """
        if df.empty:
            logger.warning("没有结果需要保存")
            return None, None
        
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        csv_file = os.path.join(output_dir, f'ma13_above_ma21_stocks_{timestamp}.csv')
        txt_file = os.path.join(output_dir, f'stock_codes_ma13_above_ma21_{timestamp}.txt')
        
        try:
            # 保存CSV文件
            df.to_csv(csv_file, index=False, encoding='utf-8-sig')
            logger.info(f"结果已保存到: {csv_file}")
            
            # 保存股票代码到TXT文件
            with open(txt_file, 'w', encoding='utf-8') as f:
                for _, row in df.iterrows():
                    f.write(f"{row['stock_code']}\n")
            logger.info(f"股票代码已保存到: {txt_file}")
            
            return csv_file, txt_file
            
        except Exception as e:
            logger.error(f"保存结果失败: {e}")
            return None, None
    
    def print_summary(self, df: pd.DataFrame):
        """打印筛选结果摘要"""
        if df.empty:
            return
        
        logger.info("\n" + "="*80)
        logger.info("筛选结果摘要")
        logger.info("="*80)
        logger.info(f"符合条件的股票数量: {len(df)}")
        logger.info(f"平均MA13-MA21间距: {df['ma13_ma21_gap'].mean():.2f}%")
        logger.info(f"最大MA13-MA21间距: {df['ma13_ma21_gap'].max():.2f}%")
        logger.info(f"最小MA13-MA21间距: {df['ma13_ma21_gap'].min():.2f}%")
        logger.info(f"平均{df['check_days'].iloc[0]}日涨幅: {df['total_change'].mean():.2f}%")
        
        logger.info(f"\n前10只MA13-MA21间距最大的股票:")
        logger.info("-"*80)
        for idx, row in df.head(10).iterrows():
            logger.info(
                f"{row['stock_code']} {row['stock_name']:8s} - "
                f"价格: {row['close_price']:7.2f} | "
                f"MA13: {row['ma13']:7.2f} | MA21: {row['ma21']:7.2f} | "
                f"间距: {row['ma13_ma21_gap']:6.2f}% | "
                f"{int(row['check_days'])}日涨幅: {row['total_change']:6.2f}%"
            )
        logger.info("="*80)
        
        # 输出逗号分隔的股票代码
        stock_codes = df['stock_code'].tolist()
        stock_codes_str = ','.join(stock_codes)
        
        logger.info("\n" + "="*80)
        logger.info("所有匹配的股票代码（逗号分隔）:")
        logger.info("="*80)
        logger.info(stock_codes_str)
        logger.info("="*80 + "\n")


def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(
        description='13日均线在21日均线之上筛选器 - 筛选短期均线强于中期均线的股票'
    )
    parser.add_argument(
        '--limit',
        type=int,
        help='限制筛选的股票数量（用于测试）'
    )
    parser.add_argument(
        '--stock',
        type=str,
        help='筛选指定股票代码'
    )
    parser.add_argument(
        '--days',
        type=int,
        default=20,
        help='检查的交易日数量（默认：20天）'
    )
    parser.add_argument(
        '--output-dir',
        type=str,
        default='.',
        help='输出文件目录（默认：当前目录）'
    )
    
    return parser.parse_args()


def main():
    """主函数"""
    args = parse_arguments()
    
    logger.info("="*80)
    logger.info("13日均线在21日均线之上筛选器")
    logger.info("筛选规则:")
    logger.info(f"1. 最近连续{args.days}个交易日，13日均线都在21日均线之上")
    logger.info("="*80 + "\n")
    
    screener = MA10AboveMA20Screener()
    
    # 如果指定了股票代码，只筛选该股票
    if args.stock:
        stock_info = screener.stock_dao.get_stock_info_by_code(args.stock)
        if not stock_info:
            logger.error(f"未找到股票: {args.stock}")
            return
        
        stock_codes = [(stock_info['stock_code'], stock_info['stock_name'])]
        logger.info(f"筛选指定股票: {args.stock} {stock_info['stock_name']}\n")
    else:
        stock_codes = None
    
    # 执行筛选
    start_time = datetime.now()
    df_results = screener.screen_stocks(stock_codes=stock_codes, limit=args.limit, check_days=args.days)
    end_time = datetime.now()
    
    # 打印结果摘要
    screener.print_summary(df_results)
    
    # 保存结果
    if not df_results.empty:
        csv_file, txt_file = screener.save_results(df_results, output_dir=args.output_dir)
        
        if csv_file and txt_file:
            logger.info(f"\n筛选用时: {(end_time - start_time).total_seconds():.2f} 秒")
            logger.info(f"详细结果请查看: {csv_file}")
            logger.info(f"股票代码列表: {txt_file}")
    
    logger.info("\n筛选完成！")


if __name__ == '__main__':
    main()

